import cv2
import numpy as np
from argparse import ArgumentParser


def add_salt_and_pepper_noise(img, low=0, high=255, low_prob=0.1, high_prob=0.1):
    result = np.zeros(img.shape, dtype=np.uint8)
    x, y = img.shape
    for i in range(x):
        for j in range(y):
            mark = np.random.choice([0, 1, 2], p=[1 - low_prob - high_prob, low_prob, high_prob])
            if mark == 0:
                result[i][j] = img[i][j]
            elif mark == 1:
                result[i][j] = low
            elif mark == 2:
                result[i][j] = high
    return result


def main(option):
    np.random.seed(option.random_seed)
    img = cv2.imread('lena_gray.bmp', cv2.IMREAD_GRAYSCALE)
    img = add_salt_and_pepper_noise(img, low=option.low, high=option.high, low_prob=option.low_prob, high_prob=option.high_prob)
    cv2.imwrite(option.output_file % (option.low, option.high, option.low_prob, option.high_prob), img)
    # cv2.namedWindow('Image')
    # cv2.imshow('Image', img)
    # cv2.waitKey()


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--random_seed', type=int, default=19950125)
    parser.add_argument('--low', type=int, default=0)
    parser.add_argument('--high', type=int, default=255)
    parser.add_argument('--low_prob', type=float, default=0.03)
    parser.add_argument('--high_prob', type=float, default=0.03)
    parser.add_argument('--output_file', type=str, default='output/noise/salt_and_pepper_low=%d_high=%d_lowprob=%.2f_highprob=%.2f.bmp')
    option = parser.parse_args()
    main(option)
